package com.sheepone.util.structure.tree.binary;

import com.sheepone.util.enums.CompareType;
import com.sheepone.util.structure.node.RedBlackTreeNode;

/**
 * @author Miss.杨
 * @description 红黑树
 * <p/>1. 任意一个节点是红色或者黑色的
 * <p/>2. 根节点是黑色的
 * <p/>3. 所有null节点视为黑色
 * <p/>4. 红色节点不能相邻
 * <p/>5. 从任意一个节点到其叶子节点的所有路径都包含相同数目的黑色节点
 * @since 2024/5/10 - 12:53
 */
public class RedBlackTree<K extends Comparable<? super K>, V> extends BinarySearchTree<K, V> {
    /**
     * 左旋，降低右边高度
     * 1. 将当前节点的右孩子的左孩子作为当前节点的右孩子
     * 2. 将当前节点作为当前节点的右孩子的左孩子（当前节点比其右孩子要小）
     *
     * @param down 需要旋转的根节点
     */
    private void leftRotate(RedBlackTreeNode<K, V> down) {
        // 需要左旋向上的节点
        RedBlackTreeNode<K, V> up = (RedBlackTreeNode<K, V>) down.getRight();
        // 需要托孤给down的节点
        RedBlackTreeNode<K, V> change = (RedBlackTreeNode<K, V>) up.getLeft();
        down.setRight(change);
        up.setLeft(down);
        updateRotatedParent(down, change, up);
    }

    /**
     * 右旋，降低左边高度
     * 1. 将当前节点的的左孩子的右孩子作为当前节点的左孩子
     * 2. 将当前节点作为当前节点的左孩子的右孩子（当前节点比其左孩子要大）
     *
     * @param down 需要右旋下降旋转的根节点
     */
    private void rightRotate(RedBlackTreeNode<K, V> down) {
        // 需要右旋向上的节点
        RedBlackTreeNode<K, V> up = (RedBlackTreeNode<K, V>) down.getLeft();
        // 需要托孤给down的节点
        RedBlackTreeNode<K, V> change = (RedBlackTreeNode<K, V>) up.getRight();
        down.setLeft(change);
        up.setRight(down);
        updateRotatedParent(down, change, up);
    }

    /**
     * 更新旋转是的父亲节点指向
     *
     * @param down   旋转中需要向下的节点
     * @param change 旋转中需要托孤的节点（up的孩子节点）
     * @param up     旋转中需要向上的节点
     */
    private void updateRotatedParent(RedBlackTreeNode<K, V> down, RedBlackTreeNode<K, V> change, RedBlackTreeNode<K, V> up) {
        // 更新change的父节点
        if (change != null) {
            change.setParent(down);
        }
        RedBlackTreeNode<K, V> parent = down.getParent();
        up.setParent(parent);
        down.setParent(up);
        // 如果parent==null说明需要更新根节点
        if (parent == null) {
            root = up;
            return;
        }
        if (parent.getLeft() == down) {
            parent.setLeft(up);
        } else {
            parent.setRight(up);
        }
    }

    @Override
    public void put(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("key must is not null");
        }
        RedBlackTreeNode<K, V> node = (RedBlackTreeNode<K, V>) root;
        RedBlackTreeNode<K, V> parent = null;
        CompareType compareType = null;
        // 找空位
        while (node != null) {
            compareType = compare(key, node.getKey());
            if (compareType == CompareType.EQUAL) {
                node.setValue(value);
                return;
            }
            parent = node;
            if (compareType == CompareType.LESS_THAN) {
                node = (RedBlackTreeNode<K, V>) node.getLeft();
            } else {
                node = (RedBlackTreeNode<K, V>) node.getRight();
            }
        }
        // 新增
        node = new RedBlackTreeNode<>(key, value);
        size++;
        if (parent != null) {
            if (compareType == CompareType.LESS_THAN) {
                parent.setLeft(node);
            } else {
                parent.setRight(node);
            }
            node.setParent(parent);
        } else {
            root = node;
        }
        fixDoubleRed(node);
    }

    /**
     * 修复红黑树红红相印的情况
     * <p>1.函数传递过来的当前节点（较低位）如果不是根节点，则一定存在父亲节点</p>
     *
     * @param lowRed 两个红色节点中较低位的那一个节点
     */
    private void fixDoubleRed(RedBlackTreeNode<K, V> lowRed) {
        // 1.根节点的颜色必须是黑色
        if (root == lowRed) {
            lowRed.setColor(RedBlackTreeNode.Color.BLACK);
        }
        // 2.当前节点的父亲是黑色的，树的红黑性不变，无需调整
        if (RedBlackTreeNode.isBlack(lowRed.getParent())) {
            return;
        }
        // 父亲节点存在，且为红色
        RedBlackTreeNode<K, V> parent = lowRed.getParent();
        // 获取其叔叔节点
        RedBlackTreeNode<K, V> uncle = lowRed.uncle();
        // 获取祖父节点
        RedBlackTreeNode<K, V> grandparent = parent.getParent();
        // 3.叔叔节点存在，且为红色
        if (RedBlackTreeNode.isRed(uncle)) {
            // 将父亲和叔叔节点都置为黑色
            parent.setColor(RedBlackTreeNode.Color.BLACK);
            uncle.setColor(RedBlackTreeNode.Color.BLACK);
            // 祖父节点置为红色
            grandparent.setColor(RedBlackTreeNode.Color.RED);
            // 将祖父节点作为新的当前节点，继续调整
            fixDoubleRed(grandparent);
            return;
        }
        // 4. 叔叔节点不存在，或者叔叔节点为黑色
        if (parent.isLeftChild()) {
            // 4.1 父亲为左孩子，当前节点也是左孩子，此时LL不平衡，需要右旋一次
            if (lowRed.isLeftChild()) {
                parent.setColor(RedBlackTreeNode.Color.BLACK);
            } else {
                // 4.2 父亲为左孩子，当前节点为右孩子，此时LR不平衡，需要先左旋，再右旋
                leftRotate(parent);
                lowRed.setColor(RedBlackTreeNode.Color.BLACK);
            }
            grandparent.setColor(RedBlackTreeNode.Color.RED);
            rightRotate(grandparent);
        } else {
            if (lowRed.isRightChild()) {
                // 4.3 父亲为右孩子，当前节点也是右孩子，此时RR不平衡，需要左旋一次
                parent.setColor(RedBlackTreeNode.Color.BLACK);
            } else {
                // 4.4 父亲为右孩子，当前节点为左孩子，此时RL不平衡，需要先右旋，再左旋
                rightRotate(parent);
                lowRed.setColor(RedBlackTreeNode.Color.BLACK);
            }
            grandparent.setColor(RedBlackTreeNode.Color.RED);
            leftRotate(grandparent);
        }
    }

    /**
     * 修复红黑树双黑节点的情况
     * <p>1. 如果当前高位节点不为根节点，那么一定存在父亲节点，所以parent!=null</p>
     * <p>2. 函数传递过来的当前节点为较高位，不为null，且当前节点的后继节点又是黑色，如果当前节点的sibling节点==null(黑色)，但是原本的树是平衡的，为了让sibling==null成立，当前节点就必须是红色，自相矛盾，所以sibling！=null</p>
     *
     * @param highBlack 两个黑色节点中较高位的那一个节点
     */
    private void fixDoubleBlack(RedBlackTreeNode<K, V> highBlack) {
        if (highBlack == root) {
            return;
        }
        RedBlackTreeNode<K, V> sibling = highBlack.sibling();
        RedBlackTreeNode<K, V> parent = highBlack.getParent();
        // 1. 被调整节点的兄弟为红色，此时两个侄子必定为黑色
        // 将其转化为(2.)或者(3.)
        if (RedBlackTreeNode.isRed(sibling)) {
            // 将兄弟节点置为黑色
            sibling.setColor(RedBlackTreeNode.Color.BLACK);
            // 将父亲置为红色
            parent.setColor(RedBlackTreeNode.Color.RED);
            // 旋转
            if (highBlack.isLeftChild()) {
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
            fixDoubleBlack(highBlack);
            return;
        }
        // 2. 被调整节点的兄弟为黑色，但是两个侄子都是黑色
        RedBlackTreeNode<K, V> left = (RedBlackTreeNode<K, V>) sibling.getLeft();
        RedBlackTreeNode<K, V> right = (RedBlackTreeNode<K, V>) sibling.getRight();
        if (RedBlackTreeNode.isBlack(left) && RedBlackTreeNode.isBlack(right)) {
            sibling.setColor(RedBlackTreeNode.Color.RED);
            if (RedBlackTreeNode.isRed(parent)) {
                parent.setColor(RedBlackTreeNode.Color.BLACK);
            } else {
                fixDoubleBlack(parent);
            }
            return;
        }
        // 3. 被调整节点的兄弟为黑色，但是两个侄子至少一个是红色
        if (sibling.isLeftChild()) {
            // 3.2 兄弟节点为左孩子，并且右侄子为红色，此时LR不平衡，需要先左旋兄弟节点（转化成3.1），再右旋父亲节点一次
            if (RedBlackTreeNode.isRed(right)) {
                leftRotate(sibling);
                // 左旋之后当前节点的兄弟节点是原先兄弟节点的右孩子
                sibling = right;
            } else {
                left.setColor(RedBlackTreeNode.Color.BLACK);
            }
            // 3.1 兄弟节点为左孩子，并且左侄子为红色，右侄子为黑色，此时LL不平衡，需要右旋父亲节点一次
            rightRotate(parent);
        } else {
            // 3.4 兄弟节点为右孩子，并且右侄子为红色，此时RL不平衡，需要先右旋兄弟节点（转化成3.3），再左旋父亲节点一次
            if (RedBlackTreeNode.isRed(left)) {
                rightRotate(sibling);
                // 左旋之后当前节点的兄弟节点是原先兄弟节点的右孩子
                sibling = left;
            }else{
                right.setColor(RedBlackTreeNode.Color.BLACK);
            }
            // 3.3 兄弟节点为右孩子，并且左侄子为红色，此时RR不平衡，需要左旋父亲节点一次
            leftRotate(parent);
        }
        // 兄弟节点最终被旋转上来顶替父亲节点的位置，所以需要将兄弟置为原先父亲节点的颜色，补充原本缺失的颜色
        sibling.setColor(parent.getColor());
        // parent最终被旋转下来，所以需要将parent置为黑色，补充当前节点被删除后缺失的黑色
        parent.setColor(RedBlackTreeNode.Color.BLACK);
    }

    @Override
    public void remove(K key) {
        RedBlackTreeNode<K, V> deleted = (RedBlackTreeNode<K, V>) getNode(key);
        if (deleted == null) {
            return;
        }
        doRemove(deleted);
    }

    /**
     * 删除节点
     * <p>1. 如果删除节点为红色，则不需要再平衡</>
     * <p>1.1 该节点无孩子，不影响红黑树特性5</p>
     * <p>1.2 该节点有孩子，因为执行的替换逻辑所以依旧不影响红黑树特性5，因为swapNode方法种只是交换节点的key与value不交换节点的颜色</p>
     * <p>2. 如果删除节点为黑色，则需要再平衡，可能破坏了红黑树特性5，需要调整</p>
     *
     * @param deleted 被删除节点
     */
    private void doRemove(RedBlackTreeNode<K, V> deleted) {
        // 找到被删除节点的替换节点，该替换节点要么是叶子节点，要么是只有一个右孩子节点
        RedBlackTreeNode<K, V> replaced = findReplacedNode(deleted);
        // 1. 被删除节点没有孩子
        RedBlackTreeNode<K, V> parent = deleted.getParent();
        if (replaced == null) {
            // 被删除节点是根节点&&没有孩子==清空树
            if (root == deleted) {
                root = null;
            } else {
                // 被删除节点是叶子节点，且是黑色，需要调整，因为null节点也是黑色节点符合双黑的情况
                if (RedBlackTreeNode.isBlack(deleted)) {
                    fixDoubleBlack(deleted);
                }
                // 删除的是叶子节点，直接置为null即可
                if (deleted.isLeftChild()) {
                    parent.setLeft(null);
                } else {
                    parent.setRight(null);
                }
                cutNode(deleted);
            }
            return;
        }
        // 2. 被删除节点有一个孩子
        if (deleted.getLeft() == null || deleted.getRight() == null) {
            // 被删除节点是根节点&&有一个红孩子==父子替换即可
            if (root == deleted) {
                swapNode(deleted, replaced);
                cutNode(replaced);
            } else {
                if (deleted.isLeftChild()) {
                    parent.setLeft(replaced);
                } else {
                    parent.setRight(replaced);
                }
                replaced.setParent(parent);
                cutNode(deleted);
                if (RedBlackTreeNode.isBlack(deleted) && RedBlackTreeNode.isBlack(replaced)) {
                    fixDoubleBlack(replaced);
                } else {
                    // 1. 如果删除的节点是红色的，那么替换节点是黑色的，此操作无影响
                    // 2. 如果删除的节点是黑色的，那么替换节点就是红色的，此时将替换节点置为黑色
                    // 即可使红黑树再次恢复平衡（补充了原本被删除而缺少的一个黑色节点）
                    replaced.setColor(RedBlackTreeNode.Color.BLACK);
                }
            }
            return;
        }
        // 3. 被删除节点有两个孩子
        // 将替换节点的key与value与被删除节点替换
        // 将被删除节点转化为第(1.)或者第(2.)种状况
        swapNode(deleted, replaced);
        doRemove(replaced);
    }

    private void cutNode(RedBlackTreeNode<K, V> node) {
        node.setLeft(null);
        node.setRight(null);
        node.setParent(null);
    }


    /**
     * 查询被删除节点的替换节点
     *
     * @param deleted 被删除的节点
     * @return 替换被删除的节点的节点
     */
    private RedBlackTreeNode<K, V> findReplacedNode(RedBlackTreeNode<K, V> deleted) {
        // 没有孩子节点
        RedBlackTreeNode<K, V> left = (RedBlackTreeNode<K, V>) deleted.getLeft();
        RedBlackTreeNode<K, V> right = (RedBlackTreeNode<K, V>) deleted.getRight();
        if (left == null && right == null) {
            return null;
        }
        // 有一个孩子节点
        if (left == null) {
            return right;
        }
        if (right == null) {
            return left;
        }
        // 有两个孩子节点，查找后继节点
        return (RedBlackTreeNode<K, V>) min(right);
    }
}
